from torch import device
from torch.cuda import is_available as is_cuda
from torch.backends.mps import is_available as is_mps
from pathlib import Path

ROOT = Path(__file__).parent.parent
DATA = ROOT / 'data'
MODEL = ROOT / 'model'

class DataPath:
    original = DATA / 'original.txt'
    data = DATA / 'data.txt'
    labeled = DATA / 'labeled.jsonl'
    preprocessed = DATA / 'preprocessed.txt'
    labels = DATA / 'labels.json'
    train = DATA / 'train.txt'
    vocab = DATA / 'vocab.txt'

class TrainArgs:
    epochs = 10
    lr = 1e-3

class ModelTrain:
    batch_size = 1
    embedding_size = 768
    hidden_size = 256
    dropout = 0.2
    train = TrainArgs()
    path_torch = MODEL / 'bilstm_crf.pth'
    path_bin = MODEL / 'bilstm_crf.bin'

class Config:
    device = device('cuda' if is_cuda() else \
                    'mps' if is_mps() else 'cpu')
    path = DataPath()
    model = ModelTrain()

conf = Config()
